#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL备份模块
提供备份和恢复功能
"""

import os
import re
import shutil
import multiprocessing
from datetime import datetime

from ...core.utils import (
    logger, 
    run_command, 
    get_input_with_default, 
    confirm_action,
    is_windows
)

class PostgresqlBackup:
    """PostgreSQL备份和恢复类"""
    
    def backup(self):
        """备份PostgreSQL数据库"""
        # 导入UI工具以提供更好的用户体验
        from ...core.ui_utils import print_title, print_info, print_success, print_error, print_warning, confirm

        print_title("备份PostgreSQL数据库")
        logger.info("开始备份PostgreSQL数据库...")
        
        # 获取指定参数
        backup_dir = get_input_with_default("输入保存备份的目录", "")
        db_name = get_input_with_default("输入要备份的数据库名称", "")
        
        if not backup_dir or not db_name:
            print_error("备份目录或数据库名称不能为空")
            return
        
        # 检查并创建备份目录
        if not os.path.exists(backup_dir):
            print_info(f"目录 {backup_dir} 不存在，正在创建...")
            try:
                os.makedirs(backup_dir)
                print_success(f"已成功创建目录: {backup_dir}")
            except PermissionError:
                print_error(f"无权创建目录: {backup_dir}")
                
                # 提供备用目录选项
                alternative_dirs = []
                
                # 添加一些常见的可能有写权限的目录
                if is_windows():
                    temp_dir = os.environ.get('TEMP', 'C:\\Temp')
                    user_profile = os.environ.get('USERPROFILE', '')
                    if user_profile:
                        user_docs = os.path.join(user_profile, 'Documents')
                        alternative_dirs.extend([temp_dir, user_docs])
                else:
                    alternative_dirs.extend([
                        '/tmp',
                        os.path.expanduser('~'),
                        '/var/tmp'
                    ])
                
                # 过滤出有写权限的目录
                writable_dirs = []
                for alt_dir in alternative_dirs:
                    if os.path.exists(alt_dir) and os.access(alt_dir, os.W_OK):
                        writable_dirs.append(alt_dir)
                
                if writable_dirs:
                    print_info("以下目录可能有写入权限:")
                    for i, dir_path in enumerate(writable_dirs, 1):
                        print_info(f"{i}. {dir_path}")
                    
                    alt_choice = get_input_with_default("请选择一个备用目录(输入编号)或按Enter继续尝试原目录", "")
                    
                    if alt_choice and alt_choice.isdigit() and 1 <= int(alt_choice) <= len(writable_dirs):
                        backup_dir = writable_dirs[int(alt_choice) - 1]
                        print_info(f"已选择备用目录: {backup_dir}")
                    else:
                        print_warning("继续使用原目录，但可能会失败")
                else:
                    print_warning("找不到可写入的备用目录")
                    return
        
        # 检查目录写权限
        if not os.access(backup_dir, os.W_OK):
            print_warning(f"目录 {backup_dir} 没有写入权限")
            
            # 在Linux系统上尝试授予权限
            if not is_windows():
                print_info("尝试授予目录写入权限...")
                try:
                    # 尝试修改目录权限
                    run_command(["sudo", "chown", "-R", f"postgres:postgres", backup_dir])
                    run_command(["sudo", "chmod", "755", backup_dir])
                    print_success("已成功授予目录写入权限")
                except Exception as e:
                    print_error(f"无法授予权限: {e}")
                    
                    # 确认是否继续尝试
                    if not confirm("无法授予目录权限，是否尝试在临时目录中备份?"):
                        return
                    
                    # 使用/tmp目录作为备用
                    tmp_backup_dir = "/tmp/pg_backup"
                    try:
                        os.makedirs(tmp_backup_dir, exist_ok=True)
                        run_command(["sudo", "chown", "postgres:postgres", tmp_backup_dir])
                        run_command(["sudo", "chmod", "777", tmp_backup_dir])
                        print_success(f"将使用临时目录: {tmp_backup_dir}")
                        backup_dir = tmp_backup_dir
                    except Exception as e:
                        print_error(f"无法创建临时备份目录: {e}")
                        return
            else:
                # Windows上的处理
                print_error("在Windows系统上，请选择一个您有写入权限的目录")
                return
        
        # 获取PostgreSQL连接信息和密码
        pg_host = get_input_with_default("输入PostgreSQL主机名", "localhost")
        pg_port = get_input_with_default("输入PostgreSQL端口号", "5432")
        pg_password = get_input_with_default("输入PostgreSQL密码", "")
        
        # 创建一个环境变量字典，包含原始环境变量和PGPASSWORD
        env_vars = os.environ.copy()
        env_vars['PGPASSWORD'] = pg_password
        
        # 生成时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 询问是使用pg_basebackup还是pg_dump
        backup_type = get_input_with_default(
            "选择备份类型 (1: pg_basebackup 物理备份, 2: pg_dump 逻辑备份)",
            "2"
        )
        
        if backup_type == "1":
            # 使用pg_basebackup进行物理备份
            target_dir = os.path.join(backup_dir, f"{db_name}_backup_{timestamp}")
            
            try:
                print_info(f"开始对 {db_name} 进行物理备份...")
                
                # 创建备份目录
                os.makedirs(target_dir, exist_ok=True)
                
                # 确保目录权限正确
                if not is_windows():
                    run_command(["sudo", "chown", "-R", "postgres:postgres", target_dir])
                
                # 运行pg_basebackup命令
                run_command([
                    "pg_basebackup", 
                    "-D", target_dir, 
                    "-Ft", 
                    "-z", 
                    "-P", 
                    "-h", pg_host,
                    "-p", pg_port,
                    "-U", "postgres"
                ], env=env_vars)  # 使用包含密码的环境变量
                
                # 压缩备份目录
                tar_file = f"{target_dir}.tar.gz"
                print_info(f"备份完成，正在压缩为: {tar_file}")
                run_command(["tar", "-czf", tar_file, "-C", backup_dir, os.path.basename(target_dir)], check=True)
                
                print_success(f"物理备份完成: {tar_file}")
            except Exception as e:
                print_error(f"物理备份失败: {e}")
                logger.error(f"物理备份失败: {e}")
                return
        else:
            # 使用pg_dump进行逻辑备份
            backup_file = os.path.join(backup_dir, f"{db_name}_backup_{timestamp}.sql")
            print_info(f"备份数据库 {db_name} 到文件: {backup_file}")
            
            try:
                # 确保目录权限正确
                if not is_windows():
                    # 确保目录对postgres用户可写
                    run_command(["sudo", "chown", "-R", "postgres:postgres", backup_dir])
                    # 确保备份目录的父目录对postgres用户可执行
                    parent_dir = os.path.dirname(backup_dir)
                    if parent_dir:
                        run_command(["sudo", "chmod", "755", parent_dir])
                
                # 运行pg_dump命令
                run_command([
                    "pg_dump", 
                    "-h", pg_host,
                    "-U", "postgres", 
                    "-p", pg_port,
                    "-f", backup_file, 
                    db_name
                ], env=env_vars)  # 使用包含密码的环境变量
                
                # 如果备份文件超过100MB，则压缩
                if os.path.exists(backup_file) and os.path.getsize(backup_file) > 100 * 1024 * 1024:  # 大于100MB时压缩
                    print_info(f"备份文件较大，正在压缩...")
                    run_command(["gzip", backup_file], check=True)
                    print_success(f"压缩完成: {backup_file}.gz")
                
                print_success(f"逻辑备份完成: {backup_file}")
            except Exception as e:
                print_error(f"逻辑备份失败: {e}")
                logger.error(f"逻辑备份失败: {e}")
                return
    
    def incremental_backup(self):
        """实现PostgreSQL增量备份(WAL归档)"""
        # 导入UI工具
        from ...core.ui_utils import print_title, print_info, print_success, print_error, print_warning, confirm, ProgressManager
        
        print_title("PostgreSQL增量备份(WAL归档方式)")
        print_info("此功能将配置WAL归档并执行基础备份，后续归档日志将自动保存")
        
        # 获取归档日志目录
        archive_dir = get_input_with_default("请输入WAL归档日志存储目录", "")
        if not archive_dir:
            print_error("归档目录不能为空")
            return
        
        # 创建归档日志目录(如果不存在)
        if not os.path.exists(archive_dir):
            try:
                print_info(f"创建归档日志目录: {archive_dir}")
                os.makedirs(archive_dir)
                # 在Linux上更改目录权限
                if not is_windows():
                    run_command(["sudo", "chown", "postgres:postgres", archive_dir])
                    run_command(["sudo", "chmod", "700", archive_dir])
            except Exception as e:
                print_error(f"创建归档日志目录失败: {e}")
                return
        
        # 获取基础备份目录
        backup_dir = get_input_with_default("请输入基础备份存储目录", "")
        if not backup_dir:
            print_error("备份目录不能为空")
            return
        
        # 创建备份目录(如果不存在)
        if not os.path.exists(backup_dir):
            try:
                print_info(f"创建基础备份目录: {backup_dir}")
                os.makedirs(backup_dir)
                # 在Linux上更改目录权限
                if not is_windows():
                    run_command(["sudo", "chown", "postgres:postgres", backup_dir])
                    run_command(["sudo", "chmod", "700", backup_dir])
            except Exception as e:
                print_error(f"创建基础备份目录失败: {e}")
                return
        
        # 获取PostgreSQL连接信息和密码
        pg_host = get_input_with_default("输入PostgreSQL主机名", "localhost")
        pg_port = get_input_with_default("输入PostgreSQL端口号", "5432")
        pg_password = get_input_with_default("输入PostgreSQL密码", "")
        
        # 创建一个环境变量字典，包含原始环境变量和PGPASSWORD
        env_vars = os.environ.copy()
        env_vars['PGPASSWORD'] = pg_password
        
        # 查找PostgreSQL配置文件路径
        try:
            result = run_command(
                ["psql", "-h", pg_host, "-p", pg_port, "-U", "postgres", "-t", "-c", "SHOW config_file"],
                capture_output=True,
                env=env_vars  # 使用包含密码的环境变量
            )
            config_file = result.stdout.strip()
            print_info(f"找到PostgreSQL配置文件: {config_file}")
        except Exception as e:
            print_error(f"无法获取PostgreSQL配置文件路径: {e}")
            return
        
        # 备份原始配置文件
        config_backup = f"{config_file}.bak"
        try:
            print_info("备份原始配置文件...")
            shutil.copy2(config_file, config_backup)
            print_success(f"配置文件已备份为: {config_backup}")
        except Exception as e:
            print_error(f"备份配置文件失败: {e}")
            return
        
        # 更新postgresql.conf配置以启用WAL归档
        try:
            print_info("开始修改PostgreSQL配置以启用WAL归档...")
            
            with open(config_file, 'r') as f:
                config_content = f.read()
            
            # 准备要更新的配置项
            archive_settings = [
                "wal_level = replica",
                "archive_mode = on",
                f"archive_command = 'test ! -f {archive_dir}/%f && cp %p {archive_dir}/%f'",
                "archive_timeout = 600"  # 10分钟强制切换一次WAL
            ]
            
            # 更新配置
            for setting in archive_settings:
                key = setting.split('=')[0].strip()
                if re.search(f"^{key}\\s*=", config_content, re.MULTILINE):
                    # 替换已存在的配置项
                    config_content = re.sub(
                        f"^{key}\\s*=.*$", 
                        setting, 
                        config_content, 
                        flags=re.MULTILINE
                    )
                else:
                    # 添加新的配置项
                    config_content += f"\n{setting}"
            
            # 写入更新后的配置
            with open(config_file, 'w') as f:
                f.write(config_content)
            
            print_success("PostgreSQL配置已更新，WAL归档已启用")
        except Exception as e:
            print_error(f"更新PostgreSQL配置失败: {e}")
            # 恢复备份的配置文件
            try:
                shutil.copy2(config_backup, config_file)
                print_warning(f"已恢复原始配置文件")
            except:
                print_error(f"无法恢复原始配置文件，请手动恢复: {config_backup}")
            return
        
        # 重新加载PostgreSQL服务以应用新配置
        try:
            print_info("重新加载PostgreSQL服务以应用配置...")
            self.reload()
        except Exception as e:
            print_error(f"重新加载PostgreSQL服务失败: {e}")
            return
        
        # 执行基础备份
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            base_backup_dir = os.path.join(backup_dir, f"base_backup_{timestamp}")
            
            print_info(f"开始执行基础备份到: {base_backup_dir}...")
            
            # 使用进度条显示备份进度
            with ProgressManager(total=100, description="执行基础备份") as progress:
                # 模拟备份进度
                progress.update(advance=10, description="准备备份...")
                
                # 开始备份
                run_command([
                    "pg_basebackup",
                    "-D", base_backup_dir,
                    "-Ft",  # tar格式
                    "-z",   # 启用压缩
                    "-P",   # 显示进度
                    "-h", pg_host,
                    "-p", pg_port,
                    "-U", "postgres"
                ], env=env_vars)  # 使用包含密码的环境变量
                
                progress.update(advance=90, description="完成基础备份")
            
            print_success(f"基础备份完成: {base_backup_dir}")
            
            # 创建备份标记文件，包含备份信息
            info_file = os.path.join(backup_dir, f"backup_info_{timestamp}.txt")
            with open(info_file, 'w') as f:
                f.write(f"备份时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"基础备份: {base_backup_dir}\n")
                f.write(f"WAL归档: {archive_dir}\n")
            
            print_info(f"备份信息已保存到: {info_file}")
            
            # 显示恢复说明
            print_title("恢复说明")
            print_info("当需要恢复数据库时，请按以下步骤操作:")
            print_info("1. 停止PostgreSQL服务")
            print_info("2. 清空数据目录")
            print_info(f"3. 解压基础备份: tar -xzf {base_backup_dir}/base.tar.gz -C /path/to/pgdata")
            print_info("4. 创建recovery.conf文件，内容如下:")
            print_info(f"   restore_command = 'cp {archive_dir}/%f %p'")
            print_info("   recovery_target_timeline = 'latest'")
            print_info("5. 启动PostgreSQL服务完成恢复")
            
        except Exception as e:
            print_error(f"执行基础备份失败: {e}")
            return
        
        print_success("增量备份设置完成。WAL日志将持续归档到指定目录，可用于时间点恢复(PITR)")
    
    def restore(self):
        """恢复PostgreSQL数据库备份"""
        restore_path = get_input_with_default("输入备份文件或目录的路径", "")
        db_name = get_input_with_default("输入要恢复的数据库名称", "")
        
        if not restore_path or not db_name:
            logger.error("备份文件路径和数据库名称不能为空")
            return
            
        # 检查备份文件/目录是否存在
        if not os.path.exists(restore_path):
            logger.error(f"备份文件或目录 {restore_path} 不存在")
            return
        
        # 获取密码并设置环境变量
        pg_password = get_input_with_default("输入PostgreSQL密码", "")
        
        # 获取PostgreSQL主机名和端口
        pg_host = get_input_with_default("输入PostgreSQL主机名", "localhost")
        pg_port = get_input_with_default("输入PostgreSQL端口", "5432")
        
        # 创建一个环境变量字典，包含原始环境变量和PGPASSWORD
        env_vars = os.environ.copy()
        env_vars['PGPASSWORD'] = pg_password
            
        # 检查数据库是否已存在
        try:
            result = run_command(
                ["psql", "-h", pg_host, "-p", pg_port, "-U", "postgres", "-lqt"], 
                check=True,
                capture_output=True,
                env=env_vars  # 使用包含密码的环境变量
            ).stdout
            db_exists = any(db_name in line for line in result.splitlines())
        except Exception as e:
            logger.error(f"检查数据库时出错: {e}")
            return
            
        # 如果数据库已存在，询问是否覆盖
        if db_exists:
            if not confirm_action(f"数据库 {db_name} 已存在。是否覆盖?", False):
                logger.info("用户选择不覆盖，退出")
                return
                
            try:
                # 先断开所有与数据库的连接
                logger.info(f"正在断开所有到数据库 {db_name} 的连接...")
                disconnect_command = f"""
                SELECT pg_terminate_backend(pg_stat_activity.pid)
                FROM pg_stat_activity
                WHERE pg_stat_activity.datname = '{db_name}'
                AND pid <> pg_backend_pid();
                """
                run_command(
                    ["psql", "-h", pg_host, "-p", pg_port, "-U", "postgres", "-c", disconnect_command], 
                    check=True,
                    env=env_vars  # 使用包含密码的环境变量
                )
                
                # 然后删除数据库
                logger.info(f"正在删除数据库 {db_name}...")
                run_command(
                    ["dropdb", "-h", pg_host, "-p", pg_port, "-U", "postgres", db_name], 
                    check=True,
                    env=env_vars  # 使用包含密码的环境变量
                )
            except Exception as e:
                logger.error(f"删除数据库时出错: {e}")
                return
                
        # 创建新的数据库
        try:
            logger.info(f"正在创建数据库 {db_name}...")
            run_command(
                ["createdb", "-h", pg_host, "-p", pg_port, "-U", "postgres", db_name], 
                check=True,
                env=env_vars  # 使用包含密码的环境变量
            )
        except Exception as e:
            logger.error(f"创建数据库时出错: {e}")
            return
            
        # 根据备份类型执行恢复
        try:
            if os.path.isdir(restore_path):
                logger.info(f"检测到目录格式备份，正在恢复数据库 {db_name}...")
                run_command(
                    ["pg_restore", "-h", pg_host, "-p", pg_port, "-U", "postgres", "-d", db_name, 
                     "-j", str(multiprocessing.cpu_count()), restore_path],
                    check=True,
                    env=env_vars  # 使用包含密码的环境变量
                )
            else:
                logger.info(f"检测到SQL文件备份，正在恢复数据库 {db_name}...")
                run_command(
                    ["psql", "-h", pg_host, "-p", pg_port, "-U", "postgres", "-d", db_name, "-f", restore_path],
                    check=True,
                    env=env_vars  # 使用包含密码的环境变量
                )
            logger.info(f"数据库 {db_name} 恢复完成")
        except Exception as e:
            logger.error(f"恢复过程中出错: {e}")
            
    def pull_backup_from_master(self):
        """从主服务器拉取备份"""
        master_ip = get_input_with_default("输入主数据库IP", "")
        master_port = get_input_with_default("输入主数据库端口", "5432")
        backup_dir = get_input_with_default("输入要保存到本地的目录", "")
        
        if not master_ip or not backup_dir:
            logger.error("主数据库IP和备份目录不能为空")
            return
        
        # 获取密码并设置环境变量
        pg_password = get_input_with_default("输入PostgreSQL密码", "")
        
        # 创建一个环境变量字典，包含原始环境变量和PGPASSWORD
        env_vars = os.environ.copy()
        env_vars['PGPASSWORD'] = pg_password
            
        # 清空指定目录
        if os.path.exists(backup_dir):
            logger.info(f"清理目录 {backup_dir}...")
            try:
                shutil.rmtree(backup_dir)
            except Exception as e:
                logger.error(f"清理目录时出错: {e}")
                return
                
        try:
            os.makedirs(backup_dir)
        except Exception as e:
            logger.error(f"创建备份目录时出错: {e}")
            return
            
        # 执行备份
        try:
            logger.info(f"正在从主库 {master_ip}:{master_port} 拉取备份到 {backup_dir}...")
            run_command([
                "pg_basebackup", "-h", master_ip, "-p", master_port, 
                "-D", backup_dir, "-Fp", "-Xs", "-P", "-U", "postgres"
            ], env=env_vars)  # 使用包含密码的环境变量
            
            # 修改文件权限和所有者
            logger.info(f"修改 {backup_dir} 中文件的所有者为 postgres:postgres，权限为 700...")
            run_command(["sudo", "chown", "-R", "postgres:postgres", backup_dir])
            run_command(["sudo", "chmod", "-R", "700", backup_dir])
            
            logger.info(f"从主服务器成功提取备份并保存到 {backup_dir}")
        except Exception as e:
            logger.error(f"从主服务器拉取备份时出错: {e}") 